# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment2/question2/main_radial.py
import matplotlib.pyplot
import numpy
import scipy

import hydrogen
import my_special


def _calculate_and_draw(subplot, title, hydrogen_atom):
    x = numpy.linspace(0, 1.5e-9, 1000)
    r = hydrogen_atom.radial(x)
    d = x ** 2 * r ** 2
    matplotlib.pyplot.subplot(*subplot)
    matplotlib.pyplot.title(title)
    matplotlib.pyplot.plot(x, d)


def main():
    h = hydrogen.Hydrogen(assoc_laguerre_function=my_special.assoc_laguerre)
    for l_ in range(0, 3):
        h.set_state(n=3, l_=l_)
        _calculate_and_draw((2, 3, l_ + 1), f"n=3 l={l_}", h)

    h = hydrogen.Hydrogen(assoc_laguerre_function=scipy.special.assoc_laguerre)
    for l_ in range(0, 3):
        h.set_state(n=3, l_=l_)
        _calculate_and_draw((2, 3, l_ + 1 + 3), "With SciPy", h)

    matplotlib.pyplot.subplots_adjust(hspace=0.35)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
